package com.ma.graph.mst;

import java.util.List;
import java.util.PriorityQueue;

public class Prim {
    //核心数据结构，存储[横切边]的优先级队列
    private PriorityQueue<int[]> pq;
    //类似visited,记录哪些节点已经成为最小生成树的一部分
    private boolean[] inMST;
    //记录最小生成树的权重和
    private int weightSum;

    // graph 是用邻接表表示的一幅图，
    // graph[s] 记录节点 s 所有相邻的边，
    // 三元组 int[]{from, to, weight} 表示一条边
    private List<int[]>[] graph;

    public Prim(List<int[]>[] graph) {
        this.graph = graph;
        //按照权值由小到大
        this.pq = new PriorityQueue<>((a, b) -> a[2] - b[2]);
        //n个节点
        int n = graph.length;
        this.inMST = new boolean[n];

        inMST[0] = true;
        cut(0);

        //不断进行切分，像最小生成树中添加边
        while (!pq.isEmpty()) {
            int[] edge = pq.poll();
            int to = edge[1];
            int weight = edge[2];
            if (inMST[to]) {
                continue;
            }
            weightSum += weight;
            inMST[to] = true;
            cut(to);
        }
    }

    //将S的横切边加入优先队列
    private void cut(int s) {
        for (int[] edge : graph[s]) {
            int to = edge[1];
            if (inMST[to]) {
                continue;
            }
            pq.offer(edge);
        }
    }

    // 最小生成树的权重和
    public int weightSum() {
        return weightSum;
    }

    // 判断最小生成树是否包含图中的所有节点
    public boolean allConnected() {
        for (int i = 0; i < inMST.length; i++) {
            if (!inMST[i]) {
                return false;
            }
        }
        return true;
    }
}